import networkx as nx
from collections import defaultdict
import logging

logger = logging.getLogger(__name__)

class GraphBuilder:
    def __init__(self):
        self.graph = nx.Graph()
        self.node_features = {}
        self.node_labels = {}

    def _add_all_nodes(self):
        for _, emp in self.employ_df.iterrows():
            node_id = f"emp_{emp['员工ID']}"
            self.graph.add_node(node_id, type='employee', original_id=emp['员工ID'],
                                composite_risk=emp['composite_risk'], position=emp['职位'],
                                department=emp['所属部门'], violation_intensity=emp['violation_intensity'])
        for _, policy in self.policy_df.iterrows():
            node_id = f"policy_{policy['保单号']}"
            self.graph.add_node(node_id, type='policy', original_id=policy['保单号'],
                                composite_risk=policy['composite_risk'], premium=policy['保费（元）'],
                                claim_count=policy['claim_count'], fraud_count=policy['fraud_count'])
        for _, inst in self.institution_df.iterrows():
            node_id = f"inst_{inst['医疗机构ID']}"
            self.graph.add_node(node_id, type='institution', original_id=inst['医疗机构ID'],
                                composite_risk=inst['composite_risk'], inst_type=inst['机构类型'],
                                level=inst['等级'], bad_record_score=inst['bad_record_score'])
        for _, claim in self.claim_df.iterrows():
            node_id = f"claim_{claim['理赔案件号']}"
            self.graph.add_node(node_id, type='claim', original_id=claim['理赔案件号'],
                                composite_risk=claim['composite_risk'], amount=claim['赔付金额（元）'],
                                is_fraud=claim['fraud_label'], delay_risk=claim.get('delay_risk', 0))
        for _, visit in self.visit_df.iterrows():
            node_id = f"visit_{visit['就诊ID']}"
            self.graph.add_node(node_id, type='visit', original_id=visit['就诊ID'],
                                composite_risk=visit['composite_risk'], total_cost=visit['总费用（元）'],
                                duration=visit.get('住院天数', 0), cost_anomaly=visit.get('cost_anomaly', 0))
        logger.info(f"添加节点完成: 员工{len(self.employ_df)}, 保单{len(self.policy_df)}, "
                    f"机构{len(self.institution_df)}, 理赔{len(self.claim_df)}, 就诊{len(self.visit_df)}")

    def _add_complete_edges(self):
        edge_counts = defaultdict(int)
        for _, policy in self.policy_df.iterrows():
            emp_node = f"emp_{policy['销售业务员ID']}"
            policy_node = f"policy_{policy['保单号']}"
            if self.graph.has_node(emp_node) and self.graph.has_node(policy_node):
                emp_risk = self.graph.nodes[emp_node].get('composite_risk', 0)
                policy_risk = self.graph.nodes[policy_node].get('composite_risk', 0)
                weight = 0.6 + (emp_risk + policy_risk) / 8.0
                self.graph.add_edge(emp_node, policy_node, relationship='sold', weight=weight)
                edge_counts['sold'] += 1
        for _, claim in self.claim_df.iterrows():
            policy_node = f"policy_{claim['关联保单号']}"
            claim_node = f"claim_{claim['理赔案件号']}"
            if self.graph.has_node(policy_node) and self.graph.has_node(claim_node):
                policy_risk = self.graph.nodes[policy_node].get('composite_risk', 0)
                claim_risk = self.graph.nodes[claim_node].get('composite_risk', 0)
                weight = 0.8 + (policy_risk + claim_risk) / 6.0
                self.graph.add_edge(policy_node, claim_node, relationship='claimed', weight=weight)
                edge_counts['claimed'] += 1
        for _, visit in self.visit_df.iterrows():
            patient_id = visit['患者ID']
            related_claims = self.claim_df[self.claim_df['申请人ID'] == patient_id]
            for _, claim in related_claims.iterrows():
                claim_node = f"claim_{claim['理赔案件号']}"
                visit_node = f"visit_{visit['就诊ID']}"
                if self.graph.has_node(claim_node) and self.graph.has_node(visit_node):
                    claim_risk = self.graph.nodes[claim_node].get('composite_risk', 0)
                    visit_risk = self.graph.nodes[visit_node].get('composite_risk', 0)
                    weight = 0.7 + (claim_risk + visit_risk) / 8.0
                    self.graph.add_edge(claim_node, visit_node, relationship='visit_claim', weight=weight)
                    edge_counts['visit_claim'] += 1
        for _, visit in self.visit_df.iterrows():
            visit_node = f"visit_{visit['就诊ID']}"
            inst_node = f"inst_{visit['医疗机构ID']}"
            if self.graph.has_node(visit_node) and self.graph.has_node(inst_node):
                visit_risk = self.graph.nodes[visit_node].get('composite_risk', 0)
                inst_risk = self.graph.nodes[inst_node].get('composite_risk', 0)
                weight = 0.9 + (visit_risk + inst_risk) / 10.0
                self.graph.add_edge(visit_node, inst_node, relationship='treated_at', weight=weight)
                edge_counts['treated_at'] += 1
        for _, claim in self.claim_df.iterrows():
            emp_node = f"emp_{claim['受理理赔员ID']}"
            claim_node = f"claim_{claim['理赔案件号']}"
            if self.graph.has_node(emp_node) and self.graph.has_node(claim_node):
                emp_risk = self.graph.nodes[emp_node].get('composite_risk', 0)
                claim_risk = self.graph.nodes[claim_node].get('composite_risk', 0)
                weight = 0.8 + (emp_risk + claim_risk) / 7.0
                self.graph.add_edge(emp_node, claim_node, relationship='processed', weight=weight)
                edge_counts['processed'] += 1
        for _, inst in self.institution_df.iterrows():
            inst_node = f"inst_{inst['医疗机构ID']}"
            inst_visits = self.visit_df[self.visit_df['医疗机构ID'] == inst['医疗机构ID']]
            for _, visit in inst_visits.iterrows():
                visit_node = f"visit_{visit['就诊ID']}"
                related_claims = self.claim_df[self.claim_df['申请人ID'] == visit['患者ID']]
                for _, claim in related_claims.iterrows():
                    claim_node = f"claim_{claim['理赔案件号']}"
                    emp_node = f"emp_{claim['受理理赔员ID']}"
                    if (self.graph.has_node(emp_node) and self.graph.has_node(inst_node) and
                            not self.graph.has_edge(emp_node, inst_node)):
                        emp_risk = self.graph.nodes[emp_node].get('composite_risk', 0)
                        inst_risk = self.graph.nodes[inst_node].get('composite_risk', 0)
                        weight = 0.5 + (emp_risk + inst_risk) / 12.0
                        self.graph.add_edge(emp_node, inst_node, relationship='indirect_emp_inst', weight=weight)
                        edge_counts['indirect_emp_inst'] += 1
        logger.info(f"完整边关系构建完成: {dict(edge_counts)}")

    def _create_complete_features_for_nodes(self):
        logger.info("创建完整节点特征...")
        feature_mapping = {}
        label_mapping = {}
        for _, emp in self.employ_df.iterrows():
            node_id = f"emp_{emp['员工ID']}"
            base_risk = emp['composite_risk']
            features = [
                base_risk / 10.0,
                emp['violation_intensity'] / 6.0,
                emp['position_risk'] / 4.0,
                emp['department_risk'] / 3.0,
                np.log1p(emp['processed_claims']) / 5.0,
                np.log1p(emp['sold_policies']) / 5.0,
                1 if '理赔' in str(emp['职位']) else 0,
                1 if '审核' in str(emp['职位']) else 0,
                1 if '管理' in str(emp['职位']) else 0,
                base_risk ** 2 / 100.0,
                np.sqrt(base_risk) if base_risk > 0 else 0,
            ]
            features = self._enhance_features(features, 'employee')
            feature_mapping[node_id] = features
            label_mapping[node_id] = 1 if base_risk > 5.0 else 0
        for _, policy in self.policy_df.iterrows():
            node_id = f"policy_{policy['保单号']}"
            base_risk = policy['composite_risk']
            features = [
                base_risk,
                policy['fraud_count'] / 2.0,
                policy['claim_count'] / 10.0,
                policy['premium_risk'],
                policy['frequency_risk'],
                policy['fraud_risk'] / 1.5,
                np.log1p(policy['claim_total']) / 10.0,
                base_risk ** 1.8,
            ]
            features = self._enhance_features(features, 'policy')
            feature_mapping[node_id] = features
            label_mapping[node_id] = 1 if base_risk > 2.0 or policy['fraud_count'] > 0 else 0
        for _, inst in self.institution_df.iterrows():
            node_id = f"inst_{inst['医疗机构ID']}"
            base_risk = inst['composite_risk']
            features = [
                base_risk / 8.0,
                inst['bad_record_score'] / 5.0,
                inst['type_risk'] / 3.0,
                inst['level_risk'] / 1.0,
                np.log1p(inst['visit_count']) / 8.0,
                np.log1p(inst['total_visit_cost']) / 12.0,
                1 if '私立' in str(inst['机构类型']) else 0,
                1 if '民营' in str(inst['机构类型']) else 0,
                base_risk ** 1.5,
            ]
            features = self._enhance_features(features, 'institution')
            feature_mapping[node_id] = features
            label_mapping[node_id] = 1 if base_risk > 4.0 else 0
        for _, claim in self.claim_df.iterrows():
            node_id = f"claim_{claim['理赔案件号']}"
            base_risk = claim['composite_risk']
            features = [
                base_risk / 6.0,
                claim['amount_risk'],
                claim['fraud_label'] * 2.5,
                claim.get('delay_risk', 0) / 2.0,
                np.log1p(claim['赔付金额（元）']) / 10.0,
                base_risk ** 2.0,
            ]
            features = self._enhance_features(features, 'claim')
            feature_mapping[node_id] = features
            label_mapping[node_id] = claim['fraud_label']
        for _, visit in self.visit_df.iterrows():
            node_id = f"visit_{visit['就诊ID']}"
            base_risk = visit['composite_risk']
            features = [
                base_risk / 4.0,
                visit['cost_risk'],
                visit['duration_risk'],
                visit.get('cost_anomaly', 0) / 1.5,
                np.log1p(visit['总费用（元）']) / 9.0,
                base_risk ** 1.6,
            ]
            features = self._enhance_features(features, 'visit')
            feature_mapping[node_id] = features
            label_mapping[node_id] = 1 if base_risk > 2.5 else 0
        self.node_features = feature_mapping
        self.node_labels = label_mapping
        labels = list(label_mapping.values())
        unique, counts = np.unique(labels, return_counts=True)
        logger.info(f"节点标签分布: {dict(zip(unique, counts))}")